"""
策略模块
先来一个简单策略
1.最开始，也就是2010-01-07日开始买入
2.以后n天为一个周期，调仓
3.假定，卖出手中所有股票
4.买入，近3天涨幅最高的n只股票，把资金平分n份买入
5.Strategy 最后给出的数据有
    1.买入卖出时间（每周调仓）
    2.买入卖出的股票（依据股票历史数据）
    3.买入卖出数量（策略给出）
"""

import pickle
from datetime import datetime
from datetime import timedelta
import stock

with open("data.pkl", "rb") as f:
    row_stock_data = pickle.load(f)

# 可以交易的日期列表
can_trade_date_list = list(row_stock_data.keys())
# 转换成dt对象，便于加减处理
can_trade_date_list_dt = [datetime.strptime(date, "%Y-%m-%d") for date in can_trade_date_list]


class Strategy(object):

    def __init__(self, name, max_num, sep, start_date, end_date):
        """
        :param name: 策略名
        :param max_num: 最大持仓数量
        :param sep: 调仓周期（天）
        :param start_date: 开始日期
        :param end_date: 结束日期
        """
        self.name = name
        self.max_num = max_num
        self.sep = sep
        self.start = start_date
        self.end = end_date

    # 买入股票的数据, 代码, 数量

    def get_buy_list(self, date):
        """
        :param date: 日期，在2010-01-07及以后
        :return:
        """
        buy_list = []
        # 为排序
        sort_dict = {}
        # 遍历股票实例
        for share in stock.stock_cls:
            sort_dict[share.code] = share.cal_n_day_return(date, 3)
        sorted_list = sorted(sort_dict.items(), key=lambda item: item[1], reverse=True)
        for i in range(0, self.max_num):
            buy_list.append(sorted_list[i][0])

        return buy_list

    @staticmethod
    def get_sell_dict(trade):
        sell_dict = {}
        # 遍历持仓中数量大于0的股票
        for item in trade.hold.index:
            if trade.hold.loc[item, '持仓数量'] > 0:
                sell_dict[item] = trade.hold.loc[item, '持仓数量']

        return sell_dict

    # 回测开始到结束之间的交易日，用来计算收益率
    def get_cal_date(self):
        cal_date_list = []
        dt_start = datetime.strptime(self.start, "%Y-%m-%d")
        dt_end = datetime.strptime(self.end, "%Y-%m-%d")
        for date in can_trade_date_list_dt:
            if dt_start <= date <= dt_end:
                cal_date_list.append(date)
        cal_date_list = [date.strftime("%Y-%m-%d") for date in cal_date_list]

        return cal_date_list

    # 获取调仓的日期
    def get_date_list(self):
        sep = self.sep
        # 空列表，用来存放交易日期
        trade_date_list = []
        # 转换为dt对象，以便于加减
        dt_start = datetime.strptime(self.start, "%Y-%m-%d")
        dt_end = datetime.strptime(self.end, "%Y-%m-%d")
        # now 是循环中的计数变量，给一个初始值，就是开始日期
        now = dt_start
        # 日期间隔（间隔多少天调仓）
        delta = timedelta(days=sep)
        while now < dt_end:  # now 在结束日期之前
            if now in can_trade_date_list_dt:  # 判断 now 在不在 可交易日期里
                trade_date_list.append(now)  # 如果在，加到我的交易列表里
                now = now + delta  # 并且，增加7天，以供下次判断
            else:
                now = now + timedelta(days=1)  # 如果不在，往后推一天，再判断
        trade_date_list = [date.strftime("%Y-%m-%d") for date in trade_date_list]
        return trade_date_list

    def cal_target_amount(self, total_money, buy_list, day):
        """
        :param total_money: 可用现金
        :param buy_list: 买入股票的列表
        :param day: 日期（买入）
        :param stock_num: 最大持仓数量
        :return:
        """
        amount_list = []
        max_num = self.max_num
        money_pre_share = total_money * 0.99 / max_num

        for share in buy_list:
            share_amount = money_pre_share // (100 * stock.Stock(share).get_open(day))
            amount_list.append(share_amount)
        return amount_list

    # 计算组合的总收益，截至某天
    @staticmethod
    def total_return(trade, day2):

        worth_1 = trade.money
        # 股票市值 + free money
        hold_info = trade.get_hold(day2)
        worth_2 = (hold_info["持仓数量"] * hold_info["现价（收盘价）"] * 100).sum() + trade.free_money

        total_return = (worth_2 - worth_1) / worth_1
        return total_return

# if __name__ == "__main__":
#     day1 = "2011-01-04"
#     day2 = "2012-01-04"
#     strategy1 = Strategy("测试", 5, 7, day1, day2)
#     c = strategy1.get_buy_list(day1)
#     b = strategy1.get_date_list()
#     d = strategy1.get_cal_date()
#     print(c,b,d,sep="\n")
